package online.heycm.platform.web.common.sse;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

/**
 * SSE 服务端，保存会话
 *
 * @author heycm
 * @since 2023/12/16 0:18
 */
@Slf4j
public class SseServer {
    private static final Map<String, SseSession> SESSION = new ConcurrentHashMap<>();

    public static SseSession connect(String id) {
        SseSession session = SESSION.get(id);
        if (session != null) {
            session.close();
        }
        session = new SseSession(id, SseSet.TIMEOUT);
        SESSION.put(id, session);
        connectCallback(session);
        // 连接成功消息
        session.send(SseSet.CONNECTED, SseSet.CONNECTED);
        return session;
    }

    public static SseSession getSession(String id) {
        return SESSION.get(id);
    }

    public static void close(String id) {
        SseSession session = getSession(id);
        if (session != null) {
            session.close();
        }
    }

    private static void remove(String id) {
        if (isConnected(id)) {
            log.info("SSE[{}] 下线", id);
            SESSION.remove(id);
        }
    }

    private static boolean isConnected(String id) {
        return SESSION.containsKey(id);
    }

    private static void connectCallback(SseSession session) {
        String id = session.getId();
        SseEmitter emitter = session.getSseEmitter();
        // 超时回调
        emitter.onTimeout(() -> {
            log.error("SSE[{}] timeout", id);
            emitter.complete();
        });
        // 连接完成时回调，发生 sseEmitter.complete() 或 sseEmitter.completeWithError() 后触发
        emitter.onCompletion(() -> {
            log.info("SSE[{}] completion", id);
            SseServer.remove(id);
        });
        // 发生异常时回调
        emitter.onError(throwable -> {
            log.error("SSE[{}] error", id, throwable);
            emitter.completeWithError(throwable);
        });
    }

}
